#!/usr/bin/python
# -*- coding: utf-8 -*-
from __future__ import annotations
import datetime, time
from typing import List, Set, Dict

from core import threadLocal
from core.clock import Clock
from core.context import Context
from core.dataClasses import R, SecurityGetter, TacticResultData
from core.enums import Period, ThreadName
from core.inventedtrader import InventedTrader
from core.tactic import Tactic
from event.event_manager import EventManager, Event
from infrastructure.util import dateutils
from infrastructure.util.dateutils import DateFormat
from core import logger


class Backtest:
    """
    回测模块
    """
    def __init__(self):
        self.event_manager = EventManager.getInstance(threadLocal.get(ThreadName.EVENT_MANAGER.value))
        self.context = Context.getInstance(threadLocal.get(ThreadName.CONTEXT.value))
        self.clock = Clock.getInstance(threadLocal.get(ThreadName.CLOCK.value))
        self.tactic: Tactic = None
        self.securities: Set[str] = []
        self.startDate: datetime.date = None

        r = R.sucess(f"初始化回测模块成功")
        logger.info(r)

    def addSecurities(self, securities: List[str]) -> Backtest:
        self.securities.extend(securities)
        r = R.sucess(f"添加回测股票代码{len(self.securities)}个成功")
        logger.info(r)
        return self

    def setTactic(self, tactic: Tactic) -> Backtest:
        self.tactic = tactic
        r = R.sucess(f"设置策略[{tactic.classInfo}]成功")
        logger.info(r)
        return self

    def setStartDate(self, date: str, dateFormat: DateFormat = DateFormat.Y_m_d) -> Backtest:
        if date is None:
            self.startDate = datetime.datetime.today()
        else:
            self.startDate = datetime.datetime.strptime(date, dateFormat.value)
        r = R.sucess(f"设置回测开始时间[{date}]成功")
        logger.info(r)
        return self

    def getStartDate(self):
        if self.startDate is None:
            self.startDate = datetime.datetime.strptime(dateutils.getLastestTradingDate(), DateFormat.Y_m_d.value)
        return self.startDate

    def _check(self):
        assert self.tactic is not None
        assert len(self.securities) > 0

    def _run(self):
        self._check()
        for security in self.securities:
            r = R.sucess(f"开始回测股票{security}")
            logger.info(r)
            self.clock.setCurrentDateTime(self.getStartDate().strftime(DateFormat.Y_m_d.value), DateFormat.Y_m_d, period=self.context.period)
            self.context.set_Security(security)
            while not self.clock.isEnd(unit=self.context.period):
                self.tactic.compute()
                self.clock.flow(unit=self.context.period)
                self.context.acquire_data()

    def start(self):
        self._run()


class BacktestMultiThread(Backtest):

    def __init__(self):
        super(BacktestMultiThread, self).__init__()
        self.getter: SecurityGetter = None

    def setGetter(self, getter: SecurityGetter):
        self.getter = getter
        return self

    def _check(self):
        assert self.tactic is not None
        assert self.getter is not None

    def _run(self):
        self._check()
        try:
            generator = self.getter()
        except:
            pass
        while(True):
            t = time.time()
            logger.info("开始回测股票, 起始时间（s）：" + str(t))
            if self.context.period == Period.minute30:
                print()
            try:
                securityInfo: TacticResultData = next(generator)
            except:
                try:
                    generator = self.getter()
                except:
                    pass
                time.sleep(1)
                continue
            if securityInfo is None:
                continue
            logger.info(f"回测股票{securityInfo.security}，得到股票代码，耗时：{str(time.time() - t)}s")
            self.clock.setCurrentDateTime(date=securityInfo.opdt.strftime(DateFormat.Y_m_d_H_M_S.value), dateFormat=DateFormat.Y_m_d_H_M_S, period=self.context.period)
            logger.info(f"回测股票{securityInfo.security}，设定开始时间成功（{self.clock.getDateTime().strftime(DateFormat.Y_m_d_H_M_S.value)}），耗时：{str(time.time() - t)}s")
            self.context.set_Security(securityInfo.security)
            logger.info(f"回测股票{securityInfo.security}，设定上下文股票代码成功，耗时：{str(time.time() - t)}s")
            while not self.clock.isEnd(unit=self.context.period):
                logger.info(f"回测股票{securityInfo.security}，判断是否是结尾时间成功，耗时：{str(time.time() - t)}s")
                if self.tactic.compute():
                    break
                logger.info(f"回测股票{securityInfo.security}，策略计算得到False，耗时：{str(time.time() - t)}s")
                self.clock.flow(unit=self.context.period)
                logger.info(f"回测股票{securityInfo.security}，时钟更新：（{self.clock.getDateTime().strftime(DateFormat.Y_m_d_H_M_S.value)}），耗时：{str(time.time() - t)}s")
                self.context.acquire_data()
                logger.info(f"回测股票{securityInfo.security}，获取时钟下最新数据，耗时：{str(time.time() - t)}s")